import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 定义颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
RED = (255, 0, 0)
GREEN = (0, 255, 0)

# 游戏设置
WIDTH, HEIGHT = 600, 400  # 游戏窗口大小
GRID_SIZE = 20  # 网格大小
FPS = 10  # 游戏帧率

# 创建游戏窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption('贪吃蛇游戏')
clock = pygame.time.Clock()

# 定义蛇类
class Snake:
    def __init__(self):
        self.positions = [(WIDTH // 2, HEIGHT // 2)]  # 蛇的初始位置（头部）
        self.direction = (GRID_SIZE, 0)  # 初始方向向右
        self.length = 1  # 初始长度
        self.score = 0  # 初始分数
    
    def get_head_position(self):
        """获取蛇头位置"""
        return self.positions[0]
    
    def update(self):
        """更新蛇的位置"""
        head = self.get_head_position()
        x, y = self.direction
        new_head = ((head[0] + x) % WIDTH, (head[1] + y) % HEIGHT)
        
        # 检查是否撞到自己
        if new_head in self.positions[1:]:
            return False  # 游戏结束
        
        self.positions.insert(0, new_head)
        if len(self.positions) > self.length:
            self.positions.pop()
        return True
    
    def grow(self):
        """蛇身体增长"""
        self.length += 1
        self.score += 10
    
    def change_direction(self, direction):
        """改变蛇的移动方向"""
        # 防止直接反向移动
        if (direction[0] * -1, direction[1] * -1) != self.direction:
            self.direction = direction
    
    def draw(self, surface):
        """绘制蛇"""
        for position in self.positions:
            rect = pygame.Rect(position[0], position[1], GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(surface, GREEN, rect)
            pygame.draw.rect(surface, BLACK, rect, 1)

# 定义食物类
class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize_position()
    
    def randomize_position(self):
        """随机生成食物位置"""
        self.position = (
            random.randint(0, (WIDTH - GRID_SIZE) // GRID_SIZE) * GRID_SIZE,
            random.randint(0, (HEIGHT - GRID_SIZE) // GRID_SIZE) * GRID_SIZE
        )
    
    def draw(self, surface):
        """绘制食物"""
        rect = pygame.Rect(self.position[0], self.position[1], GRID_SIZE, GRID_SIZE)
        pygame.draw.rect(surface, RED, rect)
        pygame.draw.rect(surface, BLACK, rect, 1)

# 绘制网格（可选）
def draw_grid(surface):
    """绘制网格线"""
    for x in range(0, WIDTH, GRID_SIZE):
        pygame.draw.line(surface, (40, 40, 40), (x, 0), (x, HEIGHT))
    for y in range(0, HEIGHT, GRID_SIZE):
        pygame.draw.line(surface, (40, 40, 40), (0, y), (WIDTH, y))

# 显示分数
def show_score(surface, score):
    """显示当前分数"""
    font = pygame.font.SysFont('arial', 20)
    score_text = font.render(f'分数: {score}', True, BLACK)
    surface.blit(score_text, (10, 10))

# 游戏结束画面
def game_over(surface, score):
    """显示游戏结束画面"""
    surface.fill(WHITE)
    font = pygame.font.SysFont('arial', 36)
    game_over_text = font.render('游戏结束!', True, BLACK)
    score_text = font.render(f'最终分数: {score}', True, BLACK)
    restart_text = font.render('按R键重新开始', True, BLACK)
    
    surface.blit(game_over_text, (WIDTH // 2 - 100, HEIGHT // 2 - 50))
    surface.blit(score_text, (WIDTH // 2 - 100, HEIGHT // 2))
    surface.blit(restart_text, (WIDTH // 2 - 100, HEIGHT // 2 + 50))
    pygame.display.update()

# 主游戏函数
def main():
    snake = Snake()
    food = Food()
    game_active = True
    
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            if event.type == pygame.KEYDOWN:
                if game_active:
                    # 控制蛇的移动方向
                    if event.key == pygame.K_UP:
                        snake.change_direction((0, -GRID_SIZE))
                    elif event.key == pygame.K_DOWN:
                        snake.change_direction((0, GRID_SIZE))
                    elif event.key == pygame.K_LEFT:
                        snake.change_direction((-GRID_SIZE, 0))
                    elif event.key == pygame.K_RIGHT:
                        snake.change_direction((GRID_SIZE, 0))
                else:
                    # 游戏结束后按R重新开始
                    if event.key == pygame.K_r:
                        main()
        
        if game_active:
            # 更新游戏状态
            if not snake.update():
                game_active = False
            
            # 检查是否吃到食物
            if snake.get_head_position() == food.position:
                snake.grow()
                food.randomize_position()
                # 确保食物不会出现在蛇身上
                while food.position in snake.positions:
                    food.randomize_position()
            
            # 绘制游戏界面
            screen.fill(WHITE)
            draw_grid(screen)
            snake.draw(screen)
            food.draw(screen)
            show_score(screen, snake.score)
            pygame.display.update()
        else:
            game_over(screen, snake.score)
        
        clock.tick(FPS)

# 启动游戏
if __name__ == "__main__":
    main()